package fem2.element;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IMatrixReadOnly;
import inf.jlinalg.IVector;
import inf.jlinalg.IVectorReadOnly;

import java.util.ArrayList;

import math2.LinAlgUtilities;
import fem2.Element;
import fem2.FieldInfo;
import fem2.IntegrationManager;
import fem2.MeshPart;
import fem2.NotImplementedError;
import fem2.enu.IntegrationType;
import fem2.enu.MassMatrixType;

/**
 * An element for Poisson heat equation
 * 
 * @author hbui
 * 
 */
public class PoissonElement extends Element {

	private double k; // thermal conductivity
	protected ArrayList<IVector> q = new ArrayList<IVector>(); // heat flux at
																// each
																// integration
																// point

	public PoissonElement(MeshPart mp, double k) {
		super(mp);
		this.k = k;
		setIntegrationType(IntegrationType.FULL_INTEGRATION);
	}

	@Override
	public int getFieldId(int idx) {
		if (idx == 0) {
			return FieldInfo.T;
		} else {
			throw new Error("Invalid input");
		}
	}

	@Override
	public IMatrixReadOnly computeKte(IMatrix kte) {
		MeshPart mp = getMeshPart();
		int ng = getNumIntegrationPoints();
		int n = mp.getBasis().getDimension();
		int dim = getMeshPart().getDimension();

		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xi = getIntegrationPoint(k);

			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ;

			double[][] gGradients = mp.gGradientsAt(xi);

			for (int i = 0; i < n; i++) {
				for (int j = 0; j < n; j++) {
					double tmp = 0.0;
					for (int l = 0; l < dim; l++) {
						tmp += gGradients[i][l] * gGradients[j][l];
					}
					kte.add(i, j, -this.k * tmp * coeff);
				}
			}
		}
		return kte;
	}

	@Override
	public IVectorReadOnly computeRie(IVector rie) {
		MeshPart mp = getMeshPart();
		int ng = getNumIntegrationPoints();
		int n = mp.getBasis().getDimension();
		int dim = getMeshPart().getDimension();

		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xi = getIntegrationPoint(k);
			IVector q = getHeatFlux(k);

			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ;

			double[][] gGradients = mp.gGradientsAt(xi);

			for (int i = 0; i < n; i++) {
				double tmp = 0.0;
				for (int l = 0; l < dim; l++) {
					tmp += gGradients[i][l] * q.get(l);
				}
				rie.add(i, tmp * coeff);
			}
		}
		return rie;
	}

	@Override
	public IMatrixReadOnly computeMe(IMatrix me, MassMatrixType massType) {
		IntegrationManager integrator = IntegrationManager.getInstance();
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();

		/*
		 * use integration scheme one higher order than integration scheme for
		 * stiffness
		 */
		int p1 = getIntegrationOrder() + 1;
		int ng = integrator.getNumIntegrationPoints(mp.getType(), p1);
		for (int k = 0; k < ng; k++) {
			double w = integrator.getIntegrationWeight(mp.getType(), p1, k);
			double[] xi = integrator.getIntegrationPoint(mp.getType(), p1, k);

			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ;

			double[] N = mp.valueAt(xi);

			for (int i = 0; i < n; i++) {
				for (int j = 0; j < n; j++) {
					for (int l = 0; l < dim; l++) {
						me.add(i * dim + l, j * dim + l, coeff * N[i] * N[j]);
					}
				}
			}
		}

		return me;
	}

	@Override
	public IVectorReadOnly computeRes(IVector res) {
		throw new NotImplementedError();
	}

	@Override
	public int getNumIntegrationPoints() {
		return IntegrationManager.getInstance().getNumIntegrationPoints(getMeshPart().getType(),
				getIntegrationOrder());
	}

	@Override
	public double[] getIntegrationPoint(int k) {
		return IntegrationManager.getInstance().getIntegrationPoint(getMeshPart().getType(),
				getIntegrationOrder(), k);
	}

	@Override
	public double getIntegrationWeight(int k) {
		return IntegrationManager.getInstance().getIntegrationWeight(getMeshPart().getType(),
				getIntegrationOrder(), k);
	}

	public IVector getHeatFlux(int k) {
		return q.get(k);
	}

	@Override
	public void Initialize() {
		/*
		 * initialize the sigma array for this element
		 */
		int ng = getNumIntegrationPoints();
		int dim = getMeshPart().getDimension();
		q.clear();
		for (int i = 0; i < ng; i++) {
			q.add(new ArrayVector(dim));
		}
	}

	@Override
	public void InitializeSolutionStep() {
		/**
		 * Calculate first heat flux for the element
		 */
		this.FinalizeNonlinearIteration();
	}

	@Override
	public void InitalizeNonlinearIteration() {
	}

	@Override
	public void FinalizeNonlinearIteration() {
		MeshPart mp = getMeshPart();
		double[] u = mp.getUHatLocal();
		int dim = mp.getDimension();
		int n = mp.getBasis().getDimension();
		/*
		 * calculate heat flux at each integration point
		 */
		for (int k = 0; k < getNumIntegrationPoints(); k++) {
			IVector q = getHeatFlux(k);
			double[] xi = getIntegrationPoint(k);
			double[][] gGradients = mp.gGradientsAt(xi);
			for (int i = 0; i < dim; i++) {
				double tmp = 0.0;
				for (int j = 0; j < n; j++) {
					tmp += u[j] * gGradients[j][i];
				}
				q.set(i, -this.k * tmp);
			}
		}
	}

	@Override
	public void FinalizeSolutionStep() {
	}

	@Override
	public void CalculateOutputData() {
	}
}
